{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Implementing a Custom Graph RAG System with Open Source SLMs and Ollama\n",
        "This notebook is a companion of chapter 13 of the \"Domain-Specific Small Language Models\" [book](https://www.manning.com/books/domain-specific-small-language-models), author Guglielmo Iozzia, [Manning Publications](https://www.manning.com/), 2025.  \n",
        "The code in this notebook is about implementing a custom Graph RAG (Retrieval Augmented Generation) system using only Small Language Models (SLMs). Hardware acceleration (GPU) is recommended.   \n",
        "More details about the code can be found in the related book's chapter."
      ],
      "metadata": {
        "id": "F8J7DW5zir2R"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Install the missing requirements in the Colab VM (only PyPDF2, Colab-xterm and Ollama Python not available by default)."
      ],
      "metadata": {
        "id": "NWo2YcYIj-hx"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b7AcqhxLep03"
      },
      "outputs": [],
      "source": [
        "!pip install PyPDF2==3.0.1 colab-xterm ollama"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Ollama setup\n",
        "Enable the Colab-xterm extension and start a terminal. Then install the Ollama server and start it by running the two commands below from the terminal:\n",
        "```\n",
        "curl -fsSL https://ollama.com/install.sh | sh\n",
        "\n",
        "ollama serve\n",
        "```\n",
        "\n"
      ],
      "metadata": {
        "id": "9pJH1Du-ObrA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%load_ext colabxterm"
      ],
      "metadata": {
        "id": "iK5SPvQMOfLc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%xterm"
      ],
      "metadata": {
        "id": "aWSPdLpEmk6O"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Finally, pull the Mistal 7B model from the Ollama model registry."
      ],
      "metadata": {
        "id": "mkhRg2_4lWCw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!ollama pull mistral"
      ],
      "metadata": {
        "id": "0uX5yztXOkel"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Indexing"
      ],
      "metadata": {
        "id": "PWnvFz6znAn3"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Upload some PDF documents."
      ],
      "metadata": {
        "id": "83jN1UOblz7Z"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!mkdir pdf_documents\n",
        "%cd pdf_documents\n",
        "!curl https://arxiv.org/pdf/2502.12923 --output arxiv_250212923.pdf\n",
        "#!curl https://arxiv.org/pdf/2502.12346 --output arxiv_250212346.pdf\n",
        "#!curl https://arxiv.org/pdf/2407.11534 --output arxiv_240711534.pdf\n",
        "%cd .."
      ],
      "metadata": {
        "id": "jPzt4fXnfRMA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's do some data cleanup. Define a custom function to remove references and appendices from the uploaded PDF."
      ],
      "metadata": {
        "id": "pwrDny6el4gh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from PyPDF2 import PdfReader, PdfWriter\n",
        "\n",
        "def remove_pdf_pages(input_pdf_path, output_pdf_path, pages_to_remove):\n",
        "    # Open the original PDF\n",
        "    reader = PdfReader(input_pdf_path)\n",
        "    writer = PdfWriter()\n",
        "\n",
        "    # Iterate through the original PDF pages\n",
        "    for page_num in range(len(reader.pages)):\n",
        "        # Add pages that are not in the pages_to_remove list\n",
        "        if page_num not in pages_to_remove:\n",
        "            writer.add_page(reader.pages[page_num])\n",
        "\n",
        "    # Write the new PDF to a file\n",
        "    with open(output_pdf_path, 'wb') as output_pdf_file:\n",
        "        writer.write(output_pdf_file)"
      ],
      "metadata": {
        "id": "v1oBLG4FxIjO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Apply the `remove_pdf_pages` function to the uploaded PDF."
      ],
      "metadata": {
        "id": "O7HBvOTtmVNT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "input_pdf_path = \"/content/pdf_documents/arxiv_250212923.pdf\"\n",
        "output_pdf_path = \"/content/pdf_documents/arxiv_250212923.pdf\"\n",
        "pages_to_remove = [6, 7, 8, 9, 10, 11]\n",
        "pages_to_remove = [x - 1 for x in pages_to_remove]\n",
        "remove_pdf_pages(input_pdf_path, output_pdf_path, pages_to_remove)"
      ],
      "metadata": {
        "id": "lYp8HbXCuBW3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define a custom function to extract text from uploaded PDF documents."
      ],
      "metadata": {
        "id": "2ALAe2udmjg8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import PyPDF2\n",
        "\n",
        "def extract_text_from_pdfs(pdf_dir, output_dir, placeholder=''):\n",
        "    \"\"\"\n",
        "    Extracts text from all PDF documents in a directory and saves them\n",
        "    as individual text files.\n",
        "\n",
        "    Args:\n",
        "        pdf_dir: The path to the directory containing PDF documents.\n",
        "        output_dir: The path to the directory where text files will be saved.\n",
        "    \"\"\"\n",
        "\n",
        "    # Create the output directory if it doesn't exist\n",
        "    os.makedirs(output_dir, exist_ok=True)\n",
        "\n",
        "    # Iterate through all files in the PDF directory\n",
        "    for filename in os.listdir(pdf_dir):\n",
        "        if filename.endswith(\".pdf\"):\n",
        "            pdf_path = os.path.join(pdf_dir, filename)\n",
        "            text_path = os.path.join(output_dir, filename[:-4] + \".txt\")  # Remove .pdf extension\n",
        "\n",
        "            # Extract text from the PDF\n",
        "            with open(pdf_path, \"rb\") as pdf_file, open(text_path, \"w\", encoding=\"utf-8\") as text_file:\n",
        "                pdf_reader = PyPDF2.PdfReader(pdf_file)\n",
        "                num_pages = len(pdf_reader.pages)\n",
        "\n",
        "                for page_num in range(num_pages):\n",
        "                    page = pdf_reader.pages[page_num]\n",
        "                    text = page.extract_text()\n",
        "                    if page_num == num_pages - 1:\n",
        "                        if placeholder in text:\n",
        "                            text = text.split(placeholder)[0]\n",
        "                    text_file.write(text)\n",
        "\n",
        "    print(\"Text extraction completed.\")"
      ],
      "metadata": {
        "id": "Lo0ykas9fUFI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Apply the `extract_text_from_pdfs` function to the uploaded PDF."
      ],
      "metadata": {
        "id": "sfEGxTjrmrcT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "pdf_dir = \"pdf_documents\"\n",
        "output_dir = \"extracted_text\"\n",
        "placeholder = 'Acknowledgments'\n",
        "extract_text_from_pdfs(pdf_dir, output_dir, placeholder)"
      ],
      "metadata": {
        "id": "J5oU_PCJqqU7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define a custom function to chunk the extracted text."
      ],
      "metadata": {
        "id": "162EPWnjmy7W"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def split_text_into_chunks(text_files_dir, chunk_size=1000):\n",
        "    \"\"\"\n",
        "    Splits multiple text files into chunks of text.\n",
        "\n",
        "    Args:\n",
        "        text_files_dir: The path to the directory containing text files.\n",
        "        chunk_size: The desired size of each chunk in characters (default: 1000).\n",
        "\n",
        "    Returns:\n",
        "        A dictionary where keys are filenames and values are lists of text chunks.\n",
        "    \"\"\"\n",
        "\n",
        "    chunks_by_file = {}\n",
        "    for filename in os.listdir(text_files_dir):\n",
        "        if filename.endswith(\".txt\"):\n",
        "            file_path = os.path.join(text_files_dir, filename)\n",
        "            with open(file_path, \"r\", encoding=\"utf-8\") as file:\n",
        "                text = file.read()\n",
        "\n",
        "            chunks = []\n",
        "            for i in range(0, len(text), chunk_size):\n",
        "                chunks.append(text[i:i + chunk_size])\n",
        "\n",
        "            chunks_by_file[filename] = chunks\n",
        "\n",
        "    return chunks_by_file"
      ],
      "metadata": {
        "id": "0dQOpWJVfbBF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Chunk the extracted text."
      ],
      "metadata": {
        "id": "borTWEVUnB3s"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "text_files_dir = \"extracted_text\"\n",
        "chunks_dict = split_text_into_chunks(text_files_dir, chunk_size=2000)"
      ],
      "metadata": {
        "id": "ksxvD8C9scxM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define the model for the model responses when extracting entities and relationships of a graph from the knowledge base."
      ],
      "metadata": {
        "id": "UDOqlzMFvC1j"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from pydantic import BaseModel\n",
        "\n",
        "class RawKnowledgeGraph(BaseModel):\n",
        "  entities: list[str]\n",
        "  relationships: list[str]"
      ],
      "metadata": {
        "id": "hrax3Bus7kLJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define a custom function to query the model hosted in the Ollama local server to extract entities and relationships from text chunks."
      ],
      "metadata": {
        "id": "XyRQ_N_4vRuW"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from ollama import chat\n",
        "from ollama import ChatResponse\n",
        "\n",
        "def extract_entities_and_relationships(chunks_dict, model_id):\n",
        "  results = {}\n",
        "  for filename, chunks in chunks_dict.items():\n",
        "        results[filename] = []\n",
        "        for chunk in chunks:\n",
        "            prompt = f\"\"\"Extract entities and relationships from this text:\\n\\n{chunk}\\n\\n\n",
        "Relationships must follow the format 'Relationship, (Entity1, Entity2)'\"\"\"\n",
        "            response: ChatResponse = chat(model=model_id, messages=[\n",
        "              {\n",
        "                'role': 'user',\n",
        "                'content': prompt,\n",
        "              },\n",
        "            ],\n",
        "            format=RawKnowledgeGraph.model_json_schema(),\n",
        "            options={\"temperature\":0}\n",
        "            )\n",
        "            results[filename].append(response['message']['content'])\n",
        "\n",
        "  return results"
      ],
      "metadata": {
        "id": "K9gX1Mhm33mC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Extract entities and relationships from the knowledge base."
      ],
      "metadata": {
        "id": "CmrjkAq0vzoF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "elements = extract_entities_and_relationships(chunks_dict, 'mistral')"
      ],
      "metadata": {
        "id": "hCmYQWwFtt5t"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "elements"
      ],
      "metadata": {
        "id": "EJ6SPO-U9zaM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define a custom function to query the model hosted in the Ollama local server to summarize graph elements."
      ],
      "metadata": {
        "id": "PagHvGKqv7zA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import json\n",
        "\n",
        "def summarize_elements(elements, model_id='mistral'):\n",
        "    \"\"\"\n",
        "    Summarizes elements using the specified Ollama model.\n",
        "\n",
        "    Args:\n",
        "        elements: A dictionary where keys are filenames and values are lists of\n",
        "                  dictionaries containing entities and relationships.\n",
        "        model_id: The ID of the Ollama model to use for summarization.\n",
        "\n",
        "    Returns:\n",
        "        A dictionary with the same structure as the input 'elements', but with\n",
        "        the values being summaries of the original entity/relationship lists.\n",
        "    \"\"\"\n",
        "    summaries = {}\n",
        "    for filename, chunks in elements.items():\n",
        "        summaries[filename] = []\n",
        "        for chunk in chunks:\n",
        "            try:\n",
        "                chunk_dict = json.loads(chunk)\n",
        "                entities = chunk_dict['entities']\n",
        "                relationships = chunk_dict['relationships']\n",
        "                prompt = f\"\"\"Summarize the following entities and relationships in the same structured format:\n",
        "Entities:\n",
        "{entities}\n",
        "\n",
        "Relationships:\n",
        "{relationships}\n",
        "                \"\"\"\n",
        "                response: ChatResponse = chat(model=model_id, messages=[\n",
        "                    {'role': 'user', 'content': prompt}\n",
        "                ],\n",
        "                format=RawKnowledgeGraph.model_json_schema(),\n",
        "                options={\"temperature\": 0}\n",
        "                )\n",
        "                summaries[filename].append(response['message']['content'])\n",
        "            except Exception as e:\n",
        "                print(f\"Error processing chunk for {filename}: {e}\")\n",
        "                summaries[filename].append({\"entities\": [], \"relationships\": []})\n",
        "\n",
        "    return summaries\n"
      ],
      "metadata": {
        "id": "Gwt2FnW8vCbA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Summarize the extracted graph elements."
      ],
      "metadata": {
        "id": "tAgts7SgwFOy"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "element_summaries = summarize_elements(elements)"
      ],
      "metadata": {
        "id": "x-f73HE44dnx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "element_summaries"
      ],
      "metadata": {
        "id": "RCBkkcvP7tRk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define a custom function to build a knowledge graph starting from entities and relationships identified by the SLM."
      ],
      "metadata": {
        "id": "pOgOPEkfwTCo"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import networkx as nx\n",
        "\n",
        "def build_graph(entities, relationships):\n",
        "  \"\"\"\n",
        "  Builds a graph from entities and relationships.\n",
        "\n",
        "  Args:\n",
        "    entities: A list of entities.\n",
        "    relationships: A list of relationships in various formats.\n",
        "                   Could be [('Relationship', 'Entity1', 'Entity2'), ...]\n",
        "                   or other structures.\n",
        "\n",
        "  Returns:\n",
        "    A NetworkX graph object.\n",
        "  \"\"\"\n",
        "\n",
        "  graph = nx.Graph()  # Create an undirected graph\n",
        "\n",
        "  # Add nodes (entities)\n",
        "  graph.add_nodes_from(entities)\n",
        "\n",
        "  # Add edges (relationships), handling different relationship formats\n",
        "  for relationship in relationships:\n",
        "    relationship_as_list = relationship.split(',')\n",
        "    # Assuming relationships are in the format ('Relationship', 'Entity1', 'Entity2')\n",
        "    # If not, adjust the logic accordingly\n",
        "    if len(relationship_as_list) >= 3:\n",
        "      relationship_name = relationship_as_list[0]\n",
        "      entity1 = relationship_as_list[1][2:]\n",
        "      entity2 = relationship_as_list[2][1:len(relationship_as_list[2])-1]\n",
        "      graph.add_edge(entity1, entity2, label=relationship_name)\n",
        "      print(f\"Adding relationship: {relationship}\")\n",
        "    else:\n",
        "      # Handle cases where the relationship is not in the expected format\n",
        "      print(f\"Skipping relationship: {relationship} - Unexpected format\")\n",
        "\n",
        "  return graph"
      ],
      "metadata": {
        "id": "dciO5Cy7v4sX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Cycle across the graph elements, build a knowledge graph for each chunk using the `build_graph` function and finally combine all the graph in a single one."
      ],
      "metadata": {
        "id": "8y6jAqLVwqB0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "combined_graph = nx.Graph()\n",
        "\n",
        "for filename, chunks_data in element_summaries.items():\n",
        "  for chunk_data in chunks_data:\n",
        "    try:\n",
        "      chunk_knowledge_graph = RawKnowledgeGraph.model_validate_json(chunk_data)\n",
        "      entities = chunk_knowledge_graph.entities\n",
        "      relationships = chunk_knowledge_graph.relationships\n",
        "\n",
        "      graph = build_graph(entities, relationships)\n",
        "      combined_graph = nx.compose(combined_graph, graph)\n",
        "\n",
        "    except Exception as e:\n",
        "      print(f\"Error processing data for {filename}: {e}\")\n",
        "\n",
        "print(\"Combined graph built successfully!\")\n"
      ],
      "metadata": {
        "id": "mVuyNt74UOvV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define a custom function to visualize a graph with Plotly, with color node points by the number of connections."
      ],
      "metadata": {
        "id": "wgU38wpEx6Er"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import plotly.graph_objects as go\n",
        "\n",
        "def visualize_graph_with_plotly(graph):\n",
        "  \"\"\"\n",
        "  Visualizes a NetworkX graph with Plotly, coloring node points\n",
        "  by the number of connections.\n",
        "\n",
        "  Args:\n",
        "    graph: A NetworkX graph object.\n",
        "  \"\"\"\n",
        "\n",
        "  pos = nx.spring_layout(graph)\n",
        "\n",
        "  edge_x = []\n",
        "  edge_y = []\n",
        "  for edge in graph.edges():\n",
        "    x0, y0 = pos[edge[0]]\n",
        "    x1, y1 = pos[edge[1]]\n",
        "    edge_x.append(x0)\n",
        "    edge_x.append(x1)\n",
        "    edge_x.append(None)\n",
        "    edge_y.append(y0)\n",
        "    edge_y.append(y1)\n",
        "    edge_y.append(None)\n",
        "\n",
        "  node_x = []\n",
        "  node_y = []\n",
        "  node_labels = []\n",
        "  for node in graph.nodes():\n",
        "    x, y = pos[node]\n",
        "    node_x.append(x)\n",
        "    node_y.append(y)\n",
        "    node_labels.append(node)\n",
        "\n",
        "  node_adjacencies = []\n",
        "  node_text = []\n",
        "  for node, adjacencies in enumerate(graph.adjacency()):\n",
        "    node_adjacencies.append(len(adjacencies[1]))\n",
        "    node_text.append(str(adjacencies[0]))\n",
        "\n",
        "  node_trace = go.Scatter(\n",
        "    x=node_x, y=node_y,\n",
        "    mode='markers',\n",
        "    hoverinfo='text',\n",
        "    marker=dict(\n",
        "      showscale=True,\n",
        "      colorscale='YlGnBu',\n",
        "      reversescale=True,\n",
        "      color=[],\n",
        "      size=10,\n",
        "      colorbar=dict(\n",
        "        thickness=15,\n",
        "        title='Node Connections',\n",
        "        xanchor='left',\n",
        "        titleside='right'\n",
        "      ),\n",
        "      line_width=2))\n",
        "\n",
        "  node_trace.marker.color = node_adjacencies\n",
        "  node_trace.text = node_text\n",
        "\n",
        "  edge_trace = go.Scatter(\n",
        "    x=edge_x, y=edge_y,\n",
        "    line=dict(width=0.5, color='#888'),\n",
        "    hoverinfo='none',\n",
        "    text=node_labels,\n",
        "    mode='lines')\n",
        "\n",
        "  fig = go.Figure(data=[edge_trace, node_trace],\n",
        "             layout=go.Layout(\n",
        "              title='<br>Chapter 13 Graph RAG example',\n",
        "              titlefont_size=16,\n",
        "              showlegend=False,\n",
        "              hovermode='closest',\n",
        "              margin=dict(b=20,l=5,r=5,t=40),\n",
        "              annotations=[ dict(\n",
        "                  text=\"\",\n",
        "                  showarrow=False,\n",
        "                  xref=\"paper\", yref=\"paper\",\n",
        "                  x=0.005, y=-0.002 ) ],\n",
        "              xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),\n",
        "              yaxis=dict(showgrid=False, zeroline=False, showticklabels=False))\n",
        "              )\n",
        "  fig.show()"
      ],
      "metadata": {
        "id": "40ym5gwTWCoU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Visualize the final knowledge graph."
      ],
      "metadata": {
        "id": "lzv5sOrWyDWq"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "visualize_graph_with_plotly(combined_graph)"
      ],
      "metadata": {
        "id": "SkJMBaUsgig9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define a custom function to save the combined graph to file in graphml format (to be easily exported and analyzed with other tools."
      ],
      "metadata": {
        "id": "4WzRg3b3yHuj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import networkx as nx\n",
        "\n",
        "def save_graph_to_graphml(graph, filepath):\n",
        "  \"\"\"Saves a NetworkX graph to a GraphML file.\n",
        "\n",
        "  Args:\n",
        "    graph: The NetworkX graph object to save.\n",
        "    filepath: The path to the output GraphML file.\n",
        "  \"\"\"\n",
        "  try:\n",
        "    nx.write_graphml(graph, filepath)\n",
        "    print(f\"Graph saved to {filepath}\")\n",
        "  except Exception as e:\n",
        "    print(f\"Error saving graph: {e}\")\n",
        "\n",
        "save_graph_to_graphml(combined_graph, \"combined_graph.graphml\")\n"
      ],
      "metadata": {
        "id": "yWCYavtK-XmM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define a custom function to implement the Leiden algorithm for community detection within a graph."
      ],
      "metadata": {
        "id": "ecMMsXeyyV-w"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import random\n",
        "\n",
        "def leiden_algorithm(graph, resolution=1.0, max_iterations=100):\n",
        "  \"\"\"\n",
        "  Implements the Leiden algorithm for community detection.\n",
        "\n",
        "  Args:\n",
        "    graph: A NetworkX graph object.\n",
        "    resolution: The resolution parameter (default: 1.0).\n",
        "    max_iterations: The maximum number of iterations (default: 100).\n",
        "\n",
        "  Returns:\n",
        "    A list of community assignments for each node.\n",
        "  \"\"\"\n",
        "  communities = {node: i for i, node in enumerate(graph.nodes)}\n",
        "  num_communities = len(communities)\n",
        "\n",
        "  for _ in range(max_iterations):\n",
        "    improved = False\n",
        "\n",
        "    for node in graph.nodes:\n",
        "      best_community = communities[node]\n",
        "      best_modularity_gain = 0\n",
        "\n",
        "      for neighbor in graph.neighbors(node):\n",
        "        neighbor_community = communities[neighbor]\n",
        "        if neighbor_community != best_community:\n",
        "          modularity_gain = calculate_modularity_gain(graph, node, neighbor_community, communities, resolution)\n",
        "          if modularity_gain > best_modularity_gain:\n",
        "            best_modularity_gain = modularity_gain\n",
        "            best_community = neighbor_community\n",
        "\n",
        "      if best_community != communities[node]:\n",
        "        communities[node] = best_community\n",
        "        improved = True\n",
        "\n",
        "    # Reassign singleton communities to a neighbor\n",
        "    for node in graph.nodes:\n",
        "      if len([n for n in graph.nodes if communities[n] == communities[node]]) == 1:\n",
        "        neighbors = list(graph.neighbors(node))\n",
        "        if neighbors:\n",
        "          communities[node] = communities[random.choice(neighbors)]\n",
        "\n",
        "    if improved:\n",
        "      communities, num_communities = aggregate_communities(graph, communities)\n",
        "    else:\n",
        "      break\n",
        "\n",
        "  return [communities[node] for node in graph.nodes]"
      ],
      "metadata": {
        "id": "VYzzl3GLlMee"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define a custom functio to calculate the modularity gain from moving a node to a community."
      ],
      "metadata": {
        "id": "a2j7bhDgyjro"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def calculate_modularity_gain(graph, node, community, communities, resolution):\n",
        "  \"\"\"\n",
        "  Calculates the modularity gain from moving a node to a community.\n",
        "  \"\"\"\n",
        "  m = graph.number_of_edges()\n",
        "  k_i = graph.degree(node)\n",
        "  k_i_in = sum(1 for neighbor in graph.neighbors(node) if communities[neighbor] == community)\n",
        "  sigma_tot = sum(graph.degree(v) for v in graph.nodes if communities[v] == community)\n",
        "\n",
        "  return (k_i_in - k_i * sigma_tot / (2 * m)) * resolution"
      ],
      "metadata": {
        "id": "NvlhPhn7lYk3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define a custom function to aggregate communities based on modularity optimization."
      ],
      "metadata": {
        "id": "C6BJQi1Uyq0a"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def aggregate_communities(graph, communities):\n",
        "  \"\"\"\n",
        "  Aggregates communities based on modularity optimization.\n",
        "  \"\"\"\n",
        "\n",
        "  # Create a mapping from original community IDs to new community IDs\n",
        "  community_mapping = {}\n",
        "  next_community_id = 0\n",
        "\n",
        "  # Assign new community IDs based on connected components\n",
        "  for component in nx.connected_components(graph):\n",
        "    for node in component:\n",
        "      if communities[node] not in community_mapping:\n",
        "        community_mapping[communities[node]] = next_community_id\n",
        "        next_community_id += 1\n",
        "\n",
        "      # Update community assignment for the node\n",
        "      communities[node] = community_mapping[communities[node]]\n",
        "\n",
        "  num_communities = next_community_id\n",
        "\n",
        "  return communities, num_communities"
      ],
      "metadata": {
        "id": "ypywrCEelb-A"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Use the Leiden algorithm to identify our graph communities."
      ],
      "metadata": {
        "id": "mMuiO9xsywpM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "community_assignments = leiden_algorithm(combined_graph, 0.5)"
      ],
      "metadata": {
        "id": "yQ84m-SpmwDu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for node, community in zip(combined_graph.nodes, community_assignments):\n",
        "  print(f\"Node {node} belongs to community {community}\")"
      ],
      "metadata": {
        "id": "Z4ddofPoljDq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Plot the identified communities."
      ],
      "metadata": {
        "id": "Gq0WXU4iy7ZG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import matplotlib.pyplot as plt\n",
        "\n",
        "node_list = list(combined_graph.nodes)\n",
        "\n",
        "node_colors = [community_assignments[node_list.index(node)] for node in combined_graph.nodes]\n",
        "\n",
        "nx.draw(combined_graph, with_labels=True, node_color=node_colors)\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "uf7h__82lmVF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define a custom function that, starting from the community list and the generated graph, uses the Mistral model hosted in the local Ollama server to generate community summaries."
      ],
      "metadata": {
        "id": "XdWUbkzozFcX"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def generate_community_summaries(community_list, graph, model_id=\"mistral\"):\n",
        "    \"\"\"\n",
        "    Generates summaries for each community in the graph using the Mistral model.\n",
        "\n",
        "    Args:\n",
        "        community_list: A list of communities, where each community is a list of node IDs.\n",
        "        graph: The NetworkX graph representing the knowledge graph.\n",
        "        model_id: The ID of the LLM model to use (default: \"mistral\").\n",
        "\n",
        "    Returns:\n",
        "        A dictionary where keys are community IDs and values are their summaries.\n",
        "    \"\"\"\n",
        "\n",
        "    community_summaries = {}\n",
        "    for community_id, community_nodes in enumerate(community_list):\n",
        "        # Extract text chunks associated with nodes in the current community\n",
        "        community_text = \"\"\n",
        "        for node in community_nodes:\n",
        "          # Assuming nodes correspond to filenames in chunks_dict\n",
        "          if node in chunks_dict:\n",
        "              for chunk in chunks_dict[node]:\n",
        "                  community_text += chunk + \"\\n\"\n",
        "\n",
        "        # Generate summary using the Mistral model\n",
        "        prompt = f\"\"\"Summarize the following text which represents a community of related documents:\n",
        "\n",
        "{community_text}\n",
        "        \"\"\"\n",
        "        response: ChatResponse = chat(model=model_id, messages=[\n",
        "            {\n",
        "                'role': 'user',\n",
        "                'content': prompt,\n",
        "            },\n",
        "        ], options={\"temperature\": 0})\n",
        "        summary = response['message']['content']\n",
        "        community_summaries[community_id] = summary\n",
        "\n",
        "    return community_summaries"
      ],
      "metadata": {
        "id": "R9BDS4l5kydi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Get a list of the unique community assignments."
      ],
      "metadata": {
        "id": "7g2mdNiCzRCK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "unique_communities = sorted(list(set(community_assignments)))"
      ],
      "metadata": {
        "id": "reJ15prTuVlc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "unique_communities = unique_communities[:10]"
      ],
      "metadata": {
        "id": "n2VhkEZyx-od"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Generate the community summaries for our graph."
      ],
      "metadata": {
        "id": "VnUN_VRlzZot"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Group nodes by community\n",
        "community_list = []\n",
        "for community_id in unique_communities:\n",
        "    community_nodes = [node for node, community in zip(graph.nodes, community_assignments) if community == community_id]\n",
        "    community_list.append(community_nodes)\n",
        "\n",
        "\n",
        "summaries = generate_community_summaries(community_list, combined_graph)\n",
        "summaries"
      ],
      "metadata": {
        "id": "msCixdoPlaPP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Querying"
      ],
      "metadata": {
        "id": "C0_bu46lm8yz"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define a custom function that generate answers to a given query from the community summaries returned by the `generate_community_summaries` function. It uses the local Ollama hosted SLM."
      ],
      "metadata": {
        "id": "nT3dkC511Mp_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def generate_answers(query, community_summaries, model_id=\"mistral\"):\n",
        "    \"\"\"\n",
        "    Generates answers to a given query based on community summaries.\n",
        "\n",
        "    Args:\n",
        "        query: The user's query.\n",
        "        community_summaries: A dictionary of community summaries.\n",
        "        model_id: The ID of the LLM model to use (default: \"mistral\").\n",
        "\n",
        "    Returns:\n",
        "        A dictionary where keys are community IDs and values are the answers\n",
        "        generated for the query based on the corresponding community summary.\n",
        "    \"\"\"\n",
        "    answers = {}\n",
        "    for community_id, summary in community_summaries.items():\n",
        "        prompt = f\"\"\"\n",
        "        Query: {query}\n",
        "\n",
        "        Use the following community summary to answer the query:\n",
        "\n",
        "        {summary}\n",
        "\n",
        "        Answer:\n",
        "        \"\"\"\n",
        "        response: ChatResponse = chat(model=model_id, messages=[\n",
        "            {\n",
        "                'role': 'user',\n",
        "                'content': prompt,\n",
        "            },\n",
        "        ], options={\"temperature\": 0})\n",
        "        answers[community_id] = response['message']['content']\n",
        "    return answers\n"
      ],
      "metadata": {
        "id": "gbbMIHxcm6re"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Generate the intermediate answers."
      ],
      "metadata": {
        "id": "0Fge2Qbk1b5n"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "query = 'What is the main topic here?'\n",
        "intermediate_answers = generate_answers(query, summaries)"
      ],
      "metadata": {
        "id": "ggNa9iukxKYK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "intermediate_answers"
      ],
      "metadata": {
        "id": "-fvf3BXCzGKI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define a custom function that generates the final answer from a user query, by combining the answers returned by `generate_answers function` into a final, concise response. It uses the local Ollama hosted SLM)."
      ],
      "metadata": {
        "id": "g9FmUfHn1iWj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def generate_final_answer(intermediate_answers, model_id=\"mistral\"):\n",
        "    \"\"\"\n",
        "    Combines intermediate answers into a final concise response.\n",
        "\n",
        "    Args:\n",
        "        intermediate_answers: A dictionary of intermediate answers from different communities.\n",
        "        model_id: The ID of the LLM model to use (default: \"mistral\").\n",
        "\n",
        "    Returns:\n",
        "        A string representing the final answer.\n",
        "    \"\"\"\n",
        "\n",
        "    # Combine intermediate answers into a single prompt\n",
        "    combined_answers_prompt = \"Combine the following answers into a single, concise response:\\n\\n\"\n",
        "    for community_id, answer in intermediate_answers.items():\n",
        "        combined_answers_prompt += f\"Community {community_id}: {answer}\\n\\n\"\n",
        "\n",
        "    # Generate the final answer using the LLM\n",
        "    response: ChatResponse = chat(model=model_id, messages=[\n",
        "        {\n",
        "            'role': 'user',\n",
        "            'content': combined_answers_prompt,\n",
        "        },\n",
        "    ], options={\"temperature\": 0})\n",
        "\n",
        "    final_answer = response['message']['content']\n",
        "    return final_answer\n"
      ],
      "metadata": {
        "id": "LOXaQsew2Dtp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Generate the final answer to the given query."
      ],
      "metadata": {
        "id": "SsYTK7MM1vQo"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "query = \"What is the main topic here?\"\n",
        "final_answer = generate_final_answer(intermediate_answers)\n",
        "final_answer"
      ],
      "metadata": {
        "id": "f_zpeU_Ixaw-"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}